#!/usr/env/bin python
"""
:mod:`pynchon.bio.plot.shell`
=============================

Printing and formatting arrays and strings to ``sys.stdout``.
"""

import sys
import numpy as np
from itertools import izip_longest
from textwrap import wrap

def make_colors(breakpoints, length=240):
    """
    Creates a linear color map given RGB "breakpoints". Returns an array of 
    colors of specified "length". A breakpoint is any RGB color-tuple. By 
    default returns an array valid for ``hijack_colors``.
    
    Arguments:
    
        - breakpoints (``sequence``) A ``sequence`` of ``tuples`` of RGB values 
          each RGB value is a floating point number from ``0.`` to ``1.``.
        - length (``int``) [default: 240] The number of rows in the returned 
          ``numpy.ndarray``.
    """
    ln = len(breakpoints)
    rng = length / (ln - 1)
    if ln < 2:
        raise ValueError('At least 2 breakpoints required.')
    breakids = []
    for i in xrange(0,ln):
        breakids.append(int(length*(i)/float(ln - 1))) # we have a dynamic range of 240

    colors = np.zeros((length, 3))
    for start in xrange(ln-1):
        lor, log, lob = breakpoints[start]
        hir, hig, hib = breakpoints[start+1]
        for id_ in xrange(*breakids[start:start+2]):
            r = (id_ % (rng)) / float(rng) # linear space
            clr = [hir*r + lor*(1-r), hig*r + log*(1-r), hib*r + lob*(1-r)]
            colors[id_] = clr
    return colors
 
def make_ruler(length, major =(20,'|'), minor=(10,'.')):
    """
    Constructs a ruler like that::
    
        |---------.---------|-------
        0                   20   
        
    Allows to specify the distance and character of minor/major ticks. Returns
    two strings for the ruler and for the numbers corresponding to major ticks.
    
    Arguments:
    
        - length (``int``) length of the constructed ruler
        - major (``tuple``) [default: (10, '.')] A tuple of the form 
          ``(N, 'X')``, where N is an integer, which means that a major-tick of 
          character 'X' will be placed every N symbols starting and including 0.
        - minor (``tuple``) [default: (10, '.')] A tuple of the form 
          ``(M, 'Y')``, where M is an integer, which means that a minor-tick of 
          character 'Y' will be placed every M symbols, but only where no 
          major-tick is placed.
    """
    ruler = []
    value = []
    for pos in range(length):
        if not pos % major[0]:
            ruler.append('|')
            s = "%" + ("-%ss" % major[0])
            value.append(s % pos)
        elif not pos % minor[0]:
            ruler.append('.')
        else:
            ruler.append('-')       
    return ("".join(ruler), "".join(value))

def hijack_colors(colors):
    """
    Sets the 240 i.e. 256 - 16 available colors of a terminal to RGB values in
    "colors". Colors should be a sequence of tuples or a rank-2 
    ``numpy.ndarray``.
    
    Arguments:
    
        - colors (``sequence`` or ``numpy.ndarray``) A sequence of 240 RGB 
          3-tuples.
    """
    if len(colors) != 240:
        raise ValueError("Array has to be of length 240.")
    init = ''
    for i, color in enumerate(colors):
        r,g,b = color
        init += "\x1b]4;%d;rgb:%02x/%02x/%02x\x1b\\" % \
                                (16+i, int(r*255),int(g*255),int(b*255))
    init += '\n'
    sys.stdout.write(init)
    
def print_array(array, colors =None):
    """
    Prints an array to ``sys.stdout``. The color map is given by "colors". 
    Custom color maps can be constructed by the ``make_colors`` function.
    
    Arguments:
    
        - array (``numpy.ndarray``) Any array which makes sense viewed as 
          ``float``.
        - colors (``sequence`` or ``numpy.ndarray``) See: ``make_colors`` and 
          ``hijack_colors``
    """
    try:
        array = np.asarray(array, dtype='float')
    except:
        raise ValueError('Array is not numeric.')
    if len(array.shape) != 2:
        raise ValueError('Expected a rank-2 array.')
    colors = colors or COLDHOT
    hijack_colors(colors)
    # shift to positive
    array -= array.min()
    # normalize to 0-1
    norm = array / array.max()
    data = []
    for row in norm:
        for el in row:
            clr_idx = int(el * 239) + 16
            assert clr_idx < 256
            data.append("\x1b[48;5;%dm  " % clr_idx)
        data.append('\x1b[0m  \n')
    sys.stdout.write("".join(data))

def print_strings(strings, raw_strings =(), ruler=True, width=80, **kwargs):
    """
    Formats and prints strings wrapped to the specified "width" and interleaved.
    Optionally a ruler is included.
    
    Arguments:
    
        - strings (``sequence``) A sequence of strings.
        - raw_strings (``sequence``) A sequence of pre-wrapped
          strings. A wrapped string is a sequence of strings of right length.
          This is useful if string includes non-printable characters.
        - ruler (``bool``) If ``True`` a ruler will be included *below* the 
          strings.
        - width (``int``) The maximum width of the print-out.
        
    Additional keyworded arguments are passed to ``make_ruler``.
    """
    maxlen = max([len(string) for string in strings])
    rul, num = make_ruler(maxlen, **kwargs)
    wvals = []
    for raw_seqs in raw_strings:
        wvals.append(raw_seqs)
    for seq in strings:
        wvals.append(wrap(seq, width=width, break_on_hyphens=False))
    if ruler:
        wvals.append(wrap(rul, width=width, break_on_hyphens=False))
        wvals.append(wrap(num, width=width, break_on_hyphens=False))
    
    for block in izip_longest(*wvals):
        for part in block:
            if part:
                sys.stdout.write(part)
            sys.stdout.write('\n')
        sys.stdout.write('\n')
        
GRAY = make_colors([
                    (0.,0.,0.),
                    (1.,1.,1.)
                    ])
COLDHOT = make_colors([
                       (0.0, 0.0, 1.0),
                       (0.0, 0.5, 1.0),
                       (1.0, 1.0, 1.0),
                       (1.0, 1.0, 0.0),
                       (1.0, 0.0, 0.0)
                       ])
COLD = make_colors([
                    (0.0, 0.0, 1.0),
                    (0.5, 0.5, 1.0),
                    (1.0, 1.0, 1.0)
                   ])
HOT = make_colors([
                   (1.0, 1.0, 1.0),
                   (1.0, 1.0, 0.0),
                   (1.0, 0.0, 0.0)
                   ])
HEAT = make_colors([
                    (0.0, 0.0, 0.0),
                    (1.0, 0.0, 0.0),
                    (1.0, 1.0, 0.0),
                    (1.0, 1.0, 1.0)
                    ])